from .resnet import *

from model.networks import define_G
from segformer.cd_transformer import Segformer
from model.SNUNet import SNUNet_ECAM


def build_model(args):
    if args.net_G == "base_transformer_pos_s4_dd8_dedim8":
        return define_G(args)
    elif args.net_G == "segformer":
        cfg_segformer = dict(
            encode_config=dict(type="mit_b4_ch4"),
            decoder_config=dict(
                in_channels=[64, 128, 320, 512],
                in_index=[0, 1, 2, 3],
                feature_strides=[4, 8, 16, 32],
                embed_dim=768,
                num_classes=2,
                dropout_ratio=0.1,
            ),
        )
        return Segformer(**cfg_segformer, args=args)
    elif args.net_G == "snunet":
        return SNUNet_ECAM(in_ch=args.in_channels, pretrained=args.pretrained)
